import time
from qturboframe.lib.webdriver.webdriver_factory import WebDriverFactory


def swith_to_new_window(timeout=10):
    '''
    切换到新弹出的窗口

    :Args:
         - timeout : 查找新窗口的次数，默认为10次

    :Returns:
        无

    :Example:
     - 例:
     @swith_to_new_window(10)
    def click_to_new_window(self):
        self.element.click()
    '''
    def wrap_func(func):
        def switch_to_new_window_decorate(self):
            driver = self._driver
            #点击按钮之前，先获取目前已经打开的窗口句柄
            globals()["OriginalHandle"] = driver.current_window_handle
            original_window_handles = driver.window_handles

            func(self)

            #切换到新弹出的窗口
            for i in range(timeout):
                window_handles = driver.window_handles
                if len(window_handles) == len(original_window_handles):
                    time.sleep(1)
                    continue

                for handle in window_handles:
                    if handle not in original_window_handles:
                        driver.switch_to.window(handle)
                break

        return switch_to_new_window_decorate
    return wrap_func

def switch_to_original_window(func):
    '''
    切换到原始窗口

    :Args:
        无

    :Returns:
        无

    :Example:
     - 例:
     @switch_to_original_window
    def click_back_to_default_window(self):
        self.element.click()
        print(self.element.text)
    '''
    def wrap_func(self):
        func(self)

        self._driver.switch_to.window(globals()["OriginalHandle"])
    return wrap_func